{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import xml.etree.ElementTree as ET\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import json\n",
    "import random\n",
    "import image\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.backends.cudnn as cudnn\n",
    "from torchvision import datasets, transforms\n",
    "from torch.autograd import Variable\n",
    "import h5py\n",
    "import numpy as np\n",
    "import cv2\n",
    "from utils import *\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def gts_boxs(frame_size,gts):\n",
    "    imgsize = (frame_size[1],frame_size[0])\n",
    "    boxs = np.zeros((len(gts),4))\n",
    "    for i in range(0,len(gts)):\n",
    "        boxs[i] = np.array(region_to_bbox(gts[i]))\n",
    "    return boxs\n",
    "def init_video(root_dataset, dataset, video):\n",
    "    video_folder = os.path.join(root_dataset, dataset, video)\n",
    "    frame_name_list = [f for f in os.listdir(video_folder) if f.endswith(\".jpg\")]\n",
    "    frame_name_list = [os.path.join(root_dataset, dataset, video, '') + s for s in frame_name_list]\n",
    "    frame_name_list.sort()\n",
    "    with Image.open(frame_name_list[0]) as img:\n",
    "        frame_sz = np.asarray(img.size)\n",
    "        frame_sz[1], frame_sz[0] = frame_sz[0], frame_sz[1]\n",
    "\n",
    "    # read the initialization from ground truth\n",
    "    gt_file = os.path.join(video_folder, 'groundtruth.txt')\n",
    "    gt = np.genfromtxt(gt_file, delimiter=',')\n",
    "    n_frames = len(frame_name_list)\n",
    "    assert n_frames == len(gt), 'Number of frames and number of GT lines should be equal.'\n",
    "\n",
    "    return gt, frame_name_list, frame_sz, n_frames\n",
    "def region_to_bbox(region, center=True):\n",
    "    n = len(region)\n",
    "    assert n==4 or n==8, ('GT region format is invalid, should have 4 or 8 entries.')\n",
    "\n",
    "    if n==4:\n",
    "        return _rect(region, center)\n",
    "    else:\n",
    "        return _poly(region, center)\n",
    "def _rect(region, center):\n",
    "    \n",
    "    if center:\n",
    "        x = region[0]\n",
    "        y = region[1]\n",
    "        w = region[2]\n",
    "        h = region[3]\n",
    "        cx = x+w/2\n",
    "        cy = y+h/2\n",
    "        return cx, cy, w, h\n",
    "    else:\n",
    "        #region[0] -= 1\n",
    "        #region[1] -= 1\n",
    "        return region\n",
    "\n",
    "\n",
    "def _poly(region, center):\n",
    "    cx = np.mean(region[::2])\n",
    "    cy = np.mean(region[1::2])\n",
    "    x1 = np.min(region[::2])\n",
    "    x2 = np.max(region[::2])\n",
    "    y1 = np.min(region[1::2])\n",
    "    y2 = np.max(region[1::2])\n",
    "    A1 = np.linalg.norm(region[0:2] - region[2:4]) * np.linalg.norm(region[2:4] - region[4:6])\n",
    "    A2 = (x2 - x1) * (y2 - y1)\n",
    "    s = np.sqrt(A1/A2)\n",
    "    w = s * (x2 - x1) + 1\n",
    "    h = s * (y2 - y1) + 1\n",
    "\n",
    "    if center:\n",
    "        return cx, cy, w, h\n",
    "    else:\n",
    "        return cx-w/2, cy-h/2, w, h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# vot2017"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "root = '/home/leeyh/Downloads/vot2017/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "names = []\n",
    "for name in glob.glob(os.path.join(root,'*')):\n",
    "    names.append(name.replace(root,''))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "root_dataset = '/home/leeyh/Downloads'\n",
    "dataset = 'vot2017'\n",
    "sequences_info = []\n",
    "for name in names:\n",
    "    gt, frame_name_list, frame_sz, n_frames = init_video(root_dataset, dataset, name)\n",
    "    gts = gts_boxs(frame_sz,gt)\n",
    "    track_ids = [[]]\n",
    "    for i in xrange(len(gts)):\n",
    "        box = gts[i]\n",
    "        track_ids[0].append([frame_name_list[i],[float(box[0]-box[2]*0.5),float(box[0]+box[2]*0.5),float(box[1]-box[3]*0.5),float(box[1]+box[3]*0.5)]])\n",
    "    sequences_info.append(track_ids)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('vot2018.txt', 'w') as outfile:\n",
    "    json.dump(sequences_info, outfile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# otb100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "root = '/home/leeyh/backup/data/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "names = []sequences_info.append(track_ids)\n",
    "for name in glob.glob(os.path.join(root,'*')):\n",
    "    names.append(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "boxes = np.loadtxt(os.path.join(names[0],'groundtruth_rect.txt'),delimiter = ',')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sequences_info = []\n",
    "for name in names:\n",
    "    track_ids = [[]]\n",
    "    try:\n",
    "        boxes = np.loadtxt(os.path.join(name,'groundtruth_rect.txt'),delimiter = ',')\n",
    "    except:\n",
    "        boxes = np.loadtxt(os.path.join(name,'groundtruth_rect.txt'))\n",
    "    img_paths = []\n",
    "    for img_path in glob.glob(os.path.join(name,'img/*')): \n",
    "        img_paths.append(img_path)\n",
    "    img_paths.sort()\n",
    "    \n",
    "    for i in xrange(len(boxes)):\n",
    "        box = boxes[i]\n",
    "        track_ids[0].append([img_paths[i],[float(box[0]-1),float(box[0]-1+box[2]),float(box[1]-1),float(box[1]-1+box[3])]])\n",
    "    sequences_info.append(track_ids)\n",
    "    \n",
    "    print name,len(img_paths),len(boxes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('otb100.txt', 'w') as outfile:\n",
    "    json.dump(sequences_info, outfile)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
